/**
 * 
 */
package timingAttack;

import java.io.Console;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.math.BigInteger;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.Arrays;
import java.util.Random;

/**
 * @author Vincent
 *
 */
public class SimpleDemoVariance 
{
	// For debugging
	static final boolean verbose = true;

	// Constants for the parameters of the demo
	static final int sampleSize = 1000;		// number of messages to try
	static final int modulusSize = 8192;	// size of modulus in bits
	static final int messageSize = 8191;	// size of messages in bits

	static BigInteger modulus;
	static BigInteger exponent; 
	static BigInteger guessForExponent;
	
	// These are the messages for which we will time modexp operations
	static BigInteger [] message = new BigInteger[sampleSize];

	// These are the total modular exponentiation times
	static long [] modExpTime; 

	// Flags for file output
	static final boolean modExpFileOutput = false;
	static final boolean bitFileOutput = false;
	static final boolean modulusFileOutput = true;

	// Filenames for file output
	static String modExpTimeFilename = modulusSize + "." + messageSize + ".modExpTime";
	static String unsetBitFilename = modulusSize + "." + messageSize + ".unsetBit";
	static String setBitFilename = modulusSize + "." + messageSize + ".setBit";
	static String modulusFilename = modulusSize + ".modulus";

	
	/**
	 * @param args
	 */
	public static void main(String[] args) 
	{
		
		// SecureRandom r = new SecureRandom();
		// exponent = new BigInteger(256,r);
		
		// Console console = System.console();
		// new BigInteger(console.readLine("Enter the exponent to guess: "));
		
		exponent = BigInteger.valueOf(199); 
		
		System.out.println("Exponent decimal representation: " + exponent.toString(10));
		System.out.println("Exponent binary representation: " + exponent.toString(2));
		System.out.println("Number of bits to guess: " + exponent.bitLength());

		// Obtain modulus
		System.out.print("Generating " + modulusSize + "-bit modulus...");
		try
		{
			modulus = generateModulus(modulusSize);
		}
		catch (Exception e)
		{
			e.printStackTrace();
			return;
		}
		
		// Save modulus to file
		if( modulusFileOutput )
		{
			PrintWriter p;
			try 
			{
				p = new PrintWriter(new FileWriter(modulusFilename));
			} 
			catch (IOException e) 
			{
				e.printStackTrace();
				return;
			}
			p.print(modulus);
			p.close();
		}
		
		// Generate timing data for messages
		System.out.print("\nGenerating timing data for " + sampleSize + " " + messageSize + "-bit messages...");
		modExpTime = timeModExp(sampleSize,messageSize,message,exponent,modulus);

		// Save total modular exponentiation times to a file
		if( modExpFileOutput )
		{
			PrintWriter p;
			try 
			{
				p = new PrintWriter(new FileWriter(modExpTimeFilename));
			} 
			catch (IOException e) 
			{
				e.printStackTrace();
				return;
			}
			writeArrayToFile(modExpTime,p);
			p.close();
		}

		// Guess the bits one at a time
		System.out.println("\n\nGuessing bits one at a time...");
		guessForExponent = guessOneBit(sampleSize);

		// Print out guess and answer
		System.out.println("\nExponent binary representation: " + exponent.toString(2));
		System.out.println("Guess binary representation: " + guessForExponent.toString(2));

		// Guess the bits two at a time
		System.out.println("\n\nGuessing bits two at a time...");
		guessForExponent = guessTwoBits(sampleSize);

		// Print out guess and answer
		System.out.println("\nExponent binary representation: " + exponent.toString(2));
		System.out.println("Guess binary representation: " + guessForExponent.toString(2));

	}

	/**
	 * Print all the elements of an array to a file
	 * 
	 * @param data
	 * @param p
	 */
	static void writeArrayToFile( long [] data, PrintWriter p )
	{
		for( int i = 0; i < data.length; i++ )
		{
			p.println(data[i]);
		}
	}
	
	/***
	 * Generate a modulus of a specified size
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeySpecException 
	 * 
	 */
	public static BigInteger generateModulus(int size) throws NoSuchAlgorithmException, InvalidKeySpecException
	{
		modulusFilename = size + ".modulus";

		// Save some time and try to read modulus from file
		try
		{
			InputStream in = new FileInputStream(modulusFilename);
			byte [] modulusBytes = new byte[in.available()];
			in.read( modulusBytes );
			in.close();
			// System.out.println("Length = " + modulusBytes.length);
			// System.out.println("Number =" + new String(modulusBytes,"UTF-8").substring(modulusBytes.length-10));
			BigInteger b = new BigInteger( new String(modulusBytes,"UTF-8") );
			System.out.print(b.bitLength() + "-bit modulus read from " + modulusFilename);
			return b;
		}
		catch(Exception e)
		{
		}
		
		// Cannot obtain modulus from file--create a new one from an RSA public key
		// http://www.javamex.com/tutorials/cryptography/rsa_encryption.shtml
		KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
		kpg.initialize(size);
		KeyPair kp = kpg.genKeyPair();
		Key publicKey = kp.getPublic();
		Key privateKey = kp.getPrivate();
		KeyFactory fact = KeyFactory.getInstance("RSA");
		RSAPublicKeySpec pub = fact.getKeySpec(kp.getPublic(), RSAPublicKeySpec.class);
		RSAPrivateKeySpec priv = fact.getKeySpec(kp.getPrivate(), RSAPrivateKeySpec.class);

		return pub.getModulus();
	}
	
	/***
	 * Generate timing data for the modular exponentiation of a
	 * specified number of messages of a given size
	 * 
	 * @param messageCount number of messages for timing
	 * @param size message size
	 * @param messageArray array to store messages in
	 * @param e exponent for modular exponentiation
	 * @param m modulus for modular exponentiation
	 * @return
	 */
	public static long [] timeModExp(int messageCount, int size, BigInteger [] messageArray, BigInteger e, BigInteger m)
	{
		long [] time = new long[messageCount];
		Random r = new SecureRandom();
		
		for( int i = 0; i < messageCount; i++ )
		{
			// Create new random messages
			messageArray[i] = new BigInteger(size,r);
			
			// Time modular exponentiation of message
			time[i] = new ModExp(m, messageArray[i], e).time; 
			
			if ((i % (messageCount/10)) == 0)
				System.out.print(i*100/messageCount + "%...");
		}
		System.out.println();
		
		return time;
	}

	/**
	 * Guess the bits of a modular exponentiation exponent using timing data
	 * Do this two bits at a time
	 * 
	 * @param messageCount number timed modular exponentation operations
	 * @return
	 */
	static BigInteger guessTwoBits(int messageCount)
	{
		BigInteger guess = BigInteger.ZERO;
		BigInteger [] runningTotal = new BigInteger[messageCount];	// of modular exp operation
		BigInteger [][] intermediateValue = 							// candidate values
			{ new BigInteger [messageCount], new BigInteger [messageCount], new BigInteger [messageCount], new BigInteger [messageCount] }; 

		long [][] remainingTime = { new long[messageCount], new long[messageCount], new long[messageCount], new long[messageCount] };
		
		// Initialize running totals, iteration times
		for( int i = 0; i < messageCount; i++ )
		{
			remainingTime[0][i] = remainingTime[1][i] = remainingTime[2][i] =remainingTime[3][i] =modExpTime[i];
			runningTotal[i] = BigInteger.ONE;
		}

		// Make a guess for every two bits of the exponent
		for( int i = 0; i < exponent.bitLength()/2; i++ )
		{
			// Open files to record times for current bit
			PrintWriter p0, p1;
			if( bitFileOutput )
			{
				try 
				{
					p0 = new PrintWriter(new FileWriter(unsetBitFilename+Integer.toString(i)));
					p1 = new PrintWriter(new FileWriter(setBitFilename+Integer.toString(i)));
				} 
				catch (IOException e) 
				{
					e.printStackTrace();
					return null;
				}
			}
			
			// Calculate times for four possibilities for each message
			for( int j = 0; j < messageCount; j++ )
			{
				ModExp modExp;
				
				// time for bits 00
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] {false, false} );
				intermediateValue[0][j] = modExp.value;
				remainingTime[0][j] -= modExp.time;
				if( bitFileOutput ) p0.println(modExp.time);
				
				// time for bits 01
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] {false, true} );
				intermediateValue[1][j] = modExp.value;
				remainingTime[1][j] -= modExp.time;
				if( bitFileOutput ) p0.println((modExp.time));

				// time for bits 10
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] {true, false} );
				intermediateValue[2][j] = modExp.value;
				remainingTime[2][j] -= modExp.time;
				if( bitFileOutput ) p1.println(modExp.time);
				
				// time for bits 11
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] {true, true} );
				intermediateValue[3][j] = modExp.value;
				remainingTime[3][j] -= modExp.time;
				if( bitFileOutput ) p1.println(modExp.time);
			}
			
			if( bitFileOutput )
			{
				p0.close();
				p1.close();
			}

			// Calculate variances for time remaining after subtracting
			// times for 0 and 1 bits
			double [] variances = { 	Statistics.variance(remainingTime[0]), Statistics.variance(remainingTime[1]),
					 					Statistics.variance(remainingTime[2]), Statistics.variance(remainingTime[3]) };

			int min = Statistics.IndexOfMin(variances);

			System.out.print( (min < 2 ? "0" : "") + Integer.toBinaryString(min));

			int [] others = new int[3];
			for( int k = 0; k < 3; k++ )
				others[k] = k + (min >= k ? 0 : 1);

			// Set the running total to the candidate total with the min variance
			// Set the candidate total with the min variance to the old running total
			// (it will be overwritten)
			BigInteger [] temp = runningTotal;
			runningTotal = intermediateValue[min];
			intermediateValue[min] = temp;

			// Set the bits in the guess if necessary
			if( (min & 2) > 0 )
				guess = guess.setBit(exponent.bitLength()-i*2-1);
			if( (min & 1) > 0 )
				guess = guess.setBit(exponent.bitLength()-i*2-2);
			
			// Output the variances
			if (verbose) System.out.println("\t" + variances[0] + " " + variances[1] + " " + variances[2] + " " + variances[3]);

			// Transfer the remaining time values for the guessed bit
			for( int k = 0; k < 3; k++ )
				remainingTime[k + (min >= k ? 0 : 1)] = Arrays.copyOf(remainingTime[min], remainingTime[min].length);
		}
		
		return guess;
	}
	
	/**
	 * Guess the bits of a modular exponentiation exponent using timing data
	 * Do this one bit at a time
	 * 
	 * @param messageCount number timed modular exponentation operations
	 * @return
	 */
	static BigInteger guessOneBit(int messageCount)
	{
		BigInteger guess = BigInteger.ZERO;
		BigInteger [] runningTotal = new BigInteger[messageCount];		// of modular exp operation
		BigInteger [][] intermediateValue = 							// candidate values
			{ new BigInteger [messageCount], new BigInteger [messageCount] }; 

		long [][] remainingTime = { new long[messageCount], new long[messageCount] };
		
		// Initialize running totals, iteration times
		for( int i = 0; i < messageCount; i++ )
		{
			remainingTime[0][i] = remainingTime[1][i] = modExpTime[i];
			runningTotal[i] = BigInteger.ONE;
		}

		// Make a guess for each bit of the exponent
		for( int i = 0; i < exponent.bitLength(); i++ )
		{
			// Open files to record times for current bit
			PrintWriter p0, p1;
			if( bitFileOutput )
			{
				try 
				{
					p0 = new PrintWriter(new FileWriter(unsetBitFilename+Integer.toString(i)));
					p1 = new PrintWriter(new FileWriter(setBitFilename+Integer.toString(i)));
				} 
				catch (IOException e) 
				{
					e.printStackTrace();
					return null;
				}
			}
			
			// Calculate times for two possibilities for each message
			for( int j = 0; j < messageCount; j++ )
			{
				ModExp modExp;
				
				// time for bit not set
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] { false } );
				intermediateValue[0][j] = modExp.value;
				remainingTime[0][j] -= modExp.time;
				if( bitFileOutput ) p0.println(modExp.time);
				
				// time for bit set
				modExp = new ModExp(runningTotal[j], modulus, message[j], new boolean [] { true } );
				intermediateValue[1][j] = modExp.value;
				remainingTime[1][j] -= modExp.time;
				if( bitFileOutput ) p1.println(modExp.time);
			}
			
			if( bitFileOutput )
			{
				p0.close();
				p1.close();
			}

			// Calculate variances for time remaining after subtracting
			// times for 0 and 1 bits
			double var0 = Statistics.variance(remainingTime[0]);
			double var1 = Statistics.variance(remainingTime[1]);

			
			int target, source;
			
			// If the Var(T) for bit=1 is smaller
			// The first bit must be 1
			if( (var0 > var1) || (i == 0) )
			{
				// bit is set
				
				// Set the running total to the candidate total for a 1-bit
				// Set the candidate total for a 1-bit to the old running total
				// (it will be overwritten)
				BigInteger [] temp = runningTotal;
				runningTotal = intermediateValue[1];
				intermediateValue[1] = temp;

				// Later copy the 1-bit remaining times to
				// the 0-bit remaining times
				target = 0;
				source = 1;
				
				// Set the bit in the guess
				guess = guess.setBit(exponent.bitLength()-i-1);
				
				System.out.print(1);
			}
			else
			{
				// bit not set

				// Set the running total to the candidate total for a 0-bit
				// Set the candidate total for a 0-bit to the old running total
				// (it will be overwritten)
				BigInteger [] temp = runningTotal;
				runningTotal = intermediateValue[0];
				intermediateValue[0] = temp;

				// Later copy the 0-bit remaining times to
				// the 1-bit remaining times
				target = 1;
				source = 0;

				System.out.print(0);
			}

			// Output the variances
			if (verbose) System.out.println("\t" + var0 + " " + var1);

			// Transfer the remaining time values for the guessed bit
			remainingTime[target] = Arrays.copyOf(remainingTime[source], remainingTime[source].length);
		}
		
		return guess;
	}

}
